from pev2.strategy.stockpool.BaseStockPool import  BaseStockPool
from pev2.factor.FactorModule import FactorModule
from pev2.data.DataModule import DataModule


class LowPeStockPool(BaseStockPool):
    """
    实现股票池选股逻辑，找到指定日期范围的候选股票
    条件：0 < PE < 30, 按从小到大排序，剔除停牌后，取前100个；再平衡周期：7个交易日
    :return: tuple，再平衡的日期列表，以及一个dict(key: 再平衡日, value: 当期的股票列表)
    """
    def getOptionStocks(self):
        factorModule = FactorModule()
        dm = DataModule()

        # 因为上证指数没有停牌不会缺数，所以用它作为交易日历
        szzsDF = dm.getKData('000001', index=True, startDate=self.startDate, endDate=self.endDate)
        allDates = list(szzsDF['date'])

        # 缓存股票和其对应有交易的日期
        codeDatesCache = dict()

        # 调整日和其对应的股票
        rebalanceDatesCodeDict = dict()
        rebalanceDates = []

        # 保存上一期的股票池
        lastPhaseCodes = []

        # 所有的交易日数
        datesCount = len(allDates)

        # 用再平衡周期作为步长循环
        for index in range(0, datesCount, self.interval):
           # 当前调整日
           rebalanceDate = allDates[index]

           # 获取本期符合条件的备选股票
           dfPE = factorModule.getSingleDateFactors('pe', rebalanceDate)
           dfPE.sort_values('pe', ascending=True, inplace=True)

           # 只保留小于30的数据
           dfPE = dfPE[(0 < dfPE['pe']) & (dfPE['pe'] < 30)]
           dfPE.set_index(['code'], inplace=True)
           thisPhaseOptionCodes = list(dfPE.index)[0:100]
           print(thisPhaseOptionCodes, flush=True)

           # 本期入选的股票代码列表
           thisPhaseCodes = []

           # 找到在上一期的股票池，但是当前停牌的股票，保留在当前股票池中
           if len(lastPhaseCodes) > 0:
               for code in lastPhaseCodes:
                   if code not in list(codeDatesCache.keys()):
                       dailyKS = dm.getKData(code, autype=None, startDate=self.startDate, endDate=self.endDate)
                       dailyKS.set_index(['date'], inplace=True)
                       codeDatesCache[code] = list(dailyKS.index)
                   if rebalanceDate not in codeDatesCache[code]:
                        thisPhaseCodes.append(code)
           print('上期停牌的股票:', flush=True)
           print(thisPhaseCodes, flush=True)

           # 剩余的位置用当前备选股票的
           optionSize = len(thisPhaseOptionCodes)
           if optionSize > (100 - len(thisPhaseCodes)):
               thisPhaseCodes += thisPhaseOptionCodes[0:100 - len(thisPhaseCodes)]
           else:
               thisPhaseCodes += thisPhaseOptionCodes

           # 当期股票池作为下次循环的上期股票池
           lastPhaseCodes = thisPhaseCodes

           # 保存到返回结果中
           rebalanceDatesCodeDict[rebalanceDate] = thisPhaseCodes
           rebalanceDates.append(rebalanceDate)

           print('当前最终的备选票：%s' % rebalanceDate, flush=True )
           print(thisPhaseCodes, flush=True)

        return rebalanceDates, rebalanceDatesCodeDict